import os
import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
from sklearn.covariance import GraphicalLasso, LedoitWolf
from sklearn.preprocessing import StandardScaler
import scipy.stats
from scipy.stats import f_oneway, kruskal
import time
import seaborn as sns

# 这里可以继续补充 mgm 分析相关函数和主流程
def preprocess_data(df, fields_core, control_vars):
    df = df[fields_core + control_vars].copy()
    # 年份字段只保留数字
    if "年份" in df.columns:
        df["年份"] = df["年份"].astype(str).str.extract(r'(\d{4})').astype(float)
    for col in df.columns:
        # 不编码年份字段
        if col != "年份" and (df[col].dtype == 'object' or df[col].dtype.name == 'category'):
            df[col] = pd.Categorical(df[col]).codes
    # 只对数值型字段填充缺失
    num_cols = df.select_dtypes(include=[np.number]).columns
    df[num_cols] = df[num_cols].fillna(df[num_cols].median())
    return df

def run_mgm(df, fields_core, control_vars, alpha=0.2, max_iter=1000):
    data = df[fields_core + control_vars].values
    data = StandardScaler().fit_transform(data)  # 新增：标准化
    model = GraphicalLasso(alpha=alpha, max_iter=max_iter)
    model.fit(data)
    adj = np.abs(model.precision_)
    return adj, fields_core + control_vars

def run_ledoitwolf(df, fields_core, control_vars):
    data = df[fields_core + control_vars].values
    model = LedoitWolf()
    model.fit(data)
    adj = np.abs(model.covariance_)
    return adj, fields_core + control_vars

def get_chinese_name(var):
    # 优先查核心变量映射，否则查控制变量映射，否则原名
    return CORE_NAME_MAP.get(var, CONTROL_NAME_MAP.get(var, var))

def build_network_from_adj(adj, var_names, df, threshold=0.10):
    G = nx.Graph()
    name_map = {v: CORE_NAME_MAP.get(v, CONTROL_NAME_MAP.get(v, v)) for v in var_names}
    for i, name_i in enumerate(var_names):
        if name_map[name_i] != "区县":  # 不添加区县节点
            G.add_node(name_map[name_i])
    for i in range(len(var_names)):
        for j in range(i+1, len(var_names)):
            x = df[var_names[i]]
            y = df[var_names[j]]
            corr, pval = scipy.stats.pearsonr(x, y)
            weight = adj[i, j]
            # 边两端都不是区县才添加
            if abs(weight) > threshold and pval < 0.05 and name_map[var_names[i]] != "区县" and name_map[var_names[j]] != "区县":
                G.add_edge(name_map[var_names[i]], name_map[var_names[j]], weight=weight, pval=pval)
    return G

